{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thorough-brake",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import minimize\n",
    "import pandas as pd\n",
    "from IPython.display import clear_output\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "\n",
    "from FBSM_library import FBSM, Objective_Functional\n",
    "from DM_library import DM\n",
    "from InitProblem_library import RKM_Dyn_System\n",
    "from Support_library import Define_Arguments, Obj_Functional_Iterations, Show_Control, Plot_Dynamics\n",
    "from Transition_Matrices_Bank import P_BC_strict_5, P_BC_strict_confidence_5, P_BC_smooth_confidence_5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f85d2f4d-1089-4e92-a3bd-4d4263213bbd",
   "metadata": {},
   "source": [
    "# M=1, m = 5 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53b3a297-bced-4a2f-9851-062beb0c2c70",
   "metadata": {},
   "source": [
    "## Set parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2854593d-2827-43a9-94b6-a20e92b006bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "M = 1  # the number of types of ordinary agents \n",
    "\n",
    "m = 5  # the number of elements in the opinion alphabet\n",
    "n = 0.6  # the fraction of ordinary agents in the system \n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the transition matrix\n",
    "\n",
    "Transition_Matrices = {}\n",
    "\n",
    "# These transition matrices outline how ordinary agents influence each other \n",
    "\n",
    "for i in range(M):\n",
    "    for j in range(M):\n",
    "        Transition_Matrices[f'{i}-{j}'] = P_BC_smooth_confidence_5\n",
    "\n",
    "# These transition matrices outline how stubborn agents influence ordinary agents \n",
    "        \n",
    "for i in range(M):\n",
    "    Transition_Matrices[f'{i}-{M}'] = P_BC_smooth_confidence_5\n",
    "    \n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the grid\n",
    "\n",
    "Tau_0 = 0\n",
    "Tau_1 = 50\n",
    "Step = 1\n",
    "Grid = np.arange(Tau_0, Tau_1+Step, Step)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Define the objective functional\n",
    "\n",
    "K = 0\n",
    "\n",
    "v = np.array([4, 3, 2, 1, 0])  # the opinion weights (terminal)\n",
    "w = K*v  # the opinion weights (integral)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting point ( y_{0} )\n",
    "\n",
    "y_0 = np.array([[0.6], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0], \n",
    "                [0]])\n",
    "\n",
    "\n",
    "n_types = [0.6]\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Create a dictionary of arguments\n",
    "\n",
    "Arguments = Define_Arguments(M, m, \n",
    "                             n, \n",
    "                             Transition_Matrices, \n",
    "                             Step, Grid, Tau_0, Tau_1, \n",
    "                             w, v, \n",
    "                             n_types)\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Initialize the starting control guess ( u_{0} (\\tau) )\n",
    "\n",
    "\n",
    "u_basic = np.array([[0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0], \n",
    "                    [0.4]])\n",
    "\n",
    "u_t_0 = []\n",
    "\n",
    "for i in range(len(Grid)-1):\n",
    "    \n",
    "    u_t_0.append(u_basic)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "208749bb-d22e-414b-9706-5e1efef0d383",
   "metadata": {},
   "source": [
    "## Run the direct method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e20ccb1-bb4f-4013-8623-d9c1473d232e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "u_t = DM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2231a51-6bd5-466f-999c-f569253df442",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e33036dd-6e1f-4577-8b13-552278ae4543",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the control function obtained\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The control function obtained:')\n",
    "\n",
    "print(Show_Control(u_t))\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "# Calculate the value of the objective functional\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t, Arguments, 1)\n",
    "\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "\n",
    "print('')\n",
    "\n",
    "print('The value of the objective functional provided by this control:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b18beb2-4dac-4210-9d01-5cac91e70543",
   "metadata": {},
   "source": [
    "## Run the FBS method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85100d47-9a74-4b94-93ce-9cd21089a43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_start = time.time()\n",
    "\n",
    "(u_t_estimations, Functional_estimations, Status) = FBSM(y_0, u_t_0, Arguments)\n",
    "\n",
    "t_end = time.time()\n",
    "\n",
    "t_elapsed = round((t_end - t_start) / 60, 1)\n",
    "\n",
    "t_elapsed_sec = round((t_end - t_start), 1)\n",
    "\n",
    "print(f'Time elapsed: {t_elapsed} min / {t_elapsed_sec} sec')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c178a467-023a-4fda-9a89-bfcb70b3dab3",
   "metadata": {},
   "source": [
    "## Show results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d10b8202-9585-4c4d-8455-dde715225e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The last approximation:')\n",
    "\n",
    "Show_Control(u_t_estimations[-1])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "print('')\n",
    "print('The best approximation:')\n",
    "\n",
    "sums = []\n",
    "\n",
    "for i in Functional_estimations:\n",
    "    sums.append(i[0]+i[1])\n",
    "\n",
    "Show_Control(u_t_estimations[np.argmin(sums)])\n",
    "\n",
    "#---------------------------------------------------------------------------\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  # last approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('The value of the objective functional provided by the last approximation:', round(value[0]+value[1], 3))\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  # the best approximation\n",
    "value = Objective_Functional(y_t, Arguments)\n",
    "print('')\n",
    "print('The value of the objective functional provided by the best approximation:', round(value[0]+value[1], 3))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "831481ab-2822-4a11-90c4-7dde39fa3662",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the dynamics of the objective functional across iterations of the FBS method\n",
    "\n",
    "FS = 15\n",
    "LW = 2\n",
    "MS = 5\n",
    "MEW = 2\n",
    "LS = 10\n",
    "\n",
    "Obj_Functional_Iterations(Functional_estimations, FS, LW, MS, MEW, LS)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the last approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[-1], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)\n",
    "\n",
    "# Plot the corresponding dynamics of state variables (the best approximation)\n",
    "\n",
    "y_t = RKM_Dyn_System(y_0, u_t_estimations[np.argmin(sums)], Arguments, 1)  \n",
    "\n",
    "LW = 1\n",
    "FS = 15\n",
    "\n",
    "Plot_Dynamics(y_t, LW, FS, Arguments)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
